#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""

"""

from pprint import pprint


class GraphNode:
    def __init__(self, label):
        self.label = label
        self.neighbors = {}     # nei_label : weight

    def add(self, nei_label, weight):
        self.neighbors[nei_label] = weight

    def sub(self, node, flow):
        self.neighbors[node] -= flow

    def __str__(self):
        return str(self.label)


class MaxFlow:
    MAX_FLOW = 1234567890

    def __init__(self, edges, source, sink):
        self.nodes = {}         # {label : node}
        self.source, self.sink = source, sink
        self.edge_flow = {}   # { edge : weight }       # NOTE: results
        self.build_adjacency(edges)
        self.flow_paths =  []   # {(each max_flow_path : flow)}

        pprint(self.nodes)
        max_flow_to_sink, max_flow_path = self._max_flow_to_sink_dijkstra()
        sum_flow_sink = 0

        while max_flow_to_sink != 0:
            print('=================== one loop =====================')
            assert max_flow_to_sink > 0
            sum_flow_sink += max_flow_to_sink

            self.flow_paths.append((max_flow_to_sink, max_flow_path))

            self._sub_path_flow(max_flow_to_sink, max_flow_path)
            max_flow_to_sink, max_flow_path = self._max_flow_to_sink_dijkstra()

        print('sum flow to sink %s' % sum_flow_sink)
        self._combine_flow_path(self.flow_paths)

        # pprint('all flow marked %s' % self.edge_flow)

    def _max_flow_to_sink_dijkstra(self):
        max_flow = {label : 0 for label in self.nodes}      # source -> 某结点max flow
        max_flow[self.source] = self.MAX_FLOW
        max_flow_path = {self.source : self.source}

        non_handled_node = set(self.nodes)

        def _find_max_node():
            m = -1
            max_node_label = None
            for node_label in non_handled_node:
                if max_flow[node_label] > m:
                    m = max_flow[node_label]
                    max_node_label = node_label

            if max_node_label:
                # delete max_node_label from non_handled_node
                non_handled_node.remove(max_node_label)

                # update max_flow by max_node_label (update max_flow_path)
                for nei_label, weight in self.nodes[max_node_label].neighbors.items():
                    cand_flow = min(m, weight)
                    if max_flow[nei_label] < cand_flow:
                        max_flow[nei_label] = cand_flow
                        max_flow_path[nei_label] = max_node_label

            return max_node_label

        max_flow_node_label = _find_max_node()

        # note: max_flow_node_label != self.sink (max flow 找至sink即可)
        while max_flow_node_label and max_flow_node_label != self.sink:
            # print(max_flow_node_label)
            max_flow_node_label = _find_max_node()

        print ('max flow map %s' )
        pprint(max_flow)
        print ('max path %s' )
        pprint(max_flow_path)
        print ('max to sink %s' % max_flow[self.sink])
        return max_flow[self.sink], max_flow_path

    def _combine_flow_path(self, flow_paths):
        for flow_cnt, path in flow_paths:
            start, end = path[self.sink], self.sink
            while True:
                if (start, end) in self.edge_flow:
                    self.edge_flow[(start, end)] += flow_cnt

                if (end, start) in self.edge_flow:
                    self.edge_flow[(end, start)] -= flow_cnt

                if start == self.source:
                    break

                start, end = path[start], start

        removes = []
        for edge, weight in self.edge_flow.items():
            if weight <= 0:
                removes.append(edge)

        for r in removes:
            del self.edge_flow[r]

    def _sub_path_flow(self, flow, max_flow_path):
        # NOTE: 并不是要sub max_flow上所有flow, 只更新source ---> sink
        def sub_edge_flow(start, end, flow):
            start_node = self.nodes[start]
            end_node = self.nodes[end]

            print('%s ---> %s' %(start, end))
            assert start_node.neighbors[end] >= flow
            start_node.neighbors[end] -= flow

            if start not in end_node.neighbors:
                end_node.add(start, flow)
            else:
                end_node.neighbors[start] += flow


        start, end = max_flow_path[self.sink], self.sink
        while True:
            sub_edge_flow(start, end, flow)
            if start == self.source:
                break
            start, end = max_flow_path[start], start

    def build_adjacency(self, edges):
        for start, end, weight in edges:
            start_node = self._get_node(start)
            end_node = self._get_node(end)
            start_node.add(end, weight)

            self.edge_flow[(start, end)] = 0

    def _get_node(self, label):
        node = self.nodes.get(label)
        if node is None:
            node = GraphNode(label)
            self.nodes[label] = node
        return node


def main():
    print("start main")
    edges = [
        ('s', 'a', 3),
        ('s', 'b', 2),

        ('a', 'b', 1),
        ('a', 'c', 3),
        ('a', 'd', 4),

        ('b', 'd', 2),

        ('c', 't', 2),
        ('d', 't', 3),
    ]

    edges_large = [
        ('s', 'a', 1000000),
        ('s', 'b', 1000000),

        ('a', 't', 1000000),
        ('b', 't', 1000000),

        ('a', 'b', 1),

    ]

    # http://www.leda-tutorial.org/en/discussion/ch05s03s04.html
    edges_example2 = [
        ('s', 'a', 90),
        ('s', 'c', 110),

        ('a', 'b', 60),
        ('a', 'd', 30),

        ('b', 'd', 40),
        ('b', 't', 50),

        ('c', 'a', 10),
        ('c', 'd', 30),

        ('d', 't', 100),
    ]

    # MIT
    edges_example3 = [
        ('s', 'a', 3),
        ('s', 'b', 2),

        ('a', 'c', 2),

        ('b', 'a', 3),
        ('b', 'd', 3),

        ('c', 'b', 1),
        ('c', 't', 3),

        ('d', 'c', 3),
        ('d', 't', 2),
    ]

    m = MaxFlow(edges, 's', 't')
    print('=================== edge flow =====================')
    pprint(m.edge_flow)
    # m = MaxFlow(edges_large, 's', 't')
    # m = MaxFlow(edges_example3, 's', 't')

if __name__ == "__main__":
    main()
